package org.systemsbiology.bigquery.gae.utils;

import com.google.api.client.extensions.appengine.http.urlfetch.UrlFetchTransport;
import com.google.api.client.googleapis.auth.oauth2.draft10.GoogleAccessProtectedResource;
import com.google.api.client.http.HttpResponseException;
import com.google.api.client.http.HttpTransport;
import com.google.api.client.json.JsonFactory;
import com.google.api.client.json.jackson.JacksonFactory;
import com.google.api.client.sample.bigquery.appengine.dashboard.AuthUtils;
import com.google.api.services.bigquery.Bigquery;
import com.google.api.services.bigquery.model.*;
import com.google.appengine.api.datastore.Entity;
import com.google.appengine.api.datastore.PreparedQuery;
import com.google.appengine.api.datastore.Query;

import java.io.IOException;
import java.util.List;
import java.util.concurrent.Callable;
import java.util.logging.Logger;

import static com.google.appengine.api.datastore.DatastoreServiceFactory.getDatastoreService;
import static com.google.common.base.Preconditions.checkNotNull;
import static java.lang.System.getProperty;
import static javax.servlet.http.HttpServletResponse.SC_INTERNAL_SERVER_ERROR;
import static javax.servlet.http.HttpServletResponse.SC_UNAUTHORIZED;
import static org.apache.commons.lang.StringUtils.isEmpty;

/**
 * @author hrovira
 */
public class BQUtils {
    private static final Logger log = Logger.getLogger(BQUtils.class.getName());

    private static final String PROJECT_ID = getProperty("com.google.api.client.sample.bigquery.appengine.dashboard.projectId");

    private final AuthUtils authUtils;
    private final GoogleAccessProtectedResource gapr;
    private final Bigquery bigquery;

    public BQUtils() throws Exception {
        String userId = lookupUserId();

        HttpTransport transport = new UrlFetchTransport();
        JsonFactory jsonFactory = new JacksonFactory();

        authUtils = new AuthUtils(userId, transport, jsonFactory);
        gapr = new GoogleAccessProtectedResource(authUtils.getAccessToken());
        bigquery = Bigquery.builder(transport, jsonFactory).setHttpRequestInitializer(gapr).build();
    }

    public Job createJob(String query) throws Exception {
        Job job = tryToDo(new JobCreator(query));
        checkNotNull(job);
        return job;
    }

    public List<JobListJobs> lookupJobs() throws Exception {
        return tryToDo(new JobsLoader()).getJobs();
    }

    public Job lookupJob(String jobId) throws Exception {
        if (isEmpty(jobId)) {
            return null;
        }

        return tryToDo(new JobLoader(jobId));
    }

    public List<DatasetListDatasets> getDatasets() throws Exception {
        DatasetList datasets = tryToDo(new Callable<DatasetList>() {
            public DatasetList call() throws IOException {
                return bigquery.datasets().list(PROJECT_ID).execute();
            }
        });
        return datasets.getDatasets();
    }

    public List<TableListTables> getTables(final String datasetId) throws Exception {
        TableList tables = tryToDo(new Callable<TableList>() {
            public TableList call() throws IOException {
                return bigquery.tables().list(PROJECT_ID, datasetId).execute();
            }
        });
        return tables.getTables();
    }

    public Table getTable(final String datasetId, final String tableId) throws Exception {
        return tryToDo(new Callable<Table>() {
            public Table call() throws IOException {
                return bigquery.tables().get(PROJECT_ID, datasetId, tableId).execute();
            }
        });
    }

    public List<Bigqueryfield> getSchemaFieldNames(Job job) throws Exception {
        if (job != null && job.getConfiguration() != null && job.getConfiguration().getQuery() != null) {
            final Tablereference tableReference = job.getConfiguration().getQuery().getDestinationTable();

            Table table = tryToDo(new Callable<Table>() {
                public Table call() throws IOException {
                    return bigquery.tables().get(tableReference.getProjectId(),
                            tableReference.getDatasetId(),
                            tableReference.getTableId()).execute();
                }
            });

            checkNotNull(table);
            checkNotNull(table.getSchema());
            checkNotNull(table.getSchema().getFields());
            return table.getSchema().getFields();
        }
        return null;
    }

    public List<Row> getTableData(Job job) throws Exception {
        if (job != null && job.getConfiguration() != null && job.getConfiguration().getQuery() != null) {
            final Tablereference tableReference =
                    job.getConfiguration().getQuery().getDestinationTable();

            TableDataList tableDataList = tryToDo(new Callable<TableDataList>() {
                public TableDataList call() throws IOException {
                    return bigquery.tabledata().list(tableReference.getProjectId(),
                            tableReference.getDatasetId(),
                            tableReference.getTableId()).execute();
                }
            });

            checkNotNull(tableDataList);
            checkNotNull(tableDataList.getRows());
            return tableDataList.getRows();
        }
        return null;
    }

    /*
     * Private Methods
     */
    private <T> T tryToDo(Callable<T> callback) throws Exception {
        int retries = 3;
        int currentTry = 0;

        Exception sdex = null;
        while (currentTry < retries) {
            currentTry++;
            try {
                return callback.call();
            } catch (Exception ex) {
                sdex = ex;
                int statusCode = getStatusCode(ex);
                log.warning(sdex.getMessage() + ":[" + statusCode + "]");
                if (statusCode == SC_UNAUTHORIZED) {
                    String access = authUtils.refreshAccessToken();
                    if (!isEmpty(access)) {
                        gapr.setAccessToken(access);
                    }
                }
            }
        }
        throw checkNotNull(sdex);
    }

    private int getStatusCode(Exception ex) {
        if (ex instanceof HttpResponseException) {
            HttpResponseException hrex = (HttpResponseException) ex;
            return hrex.getResponse().getStatusCode();
        }
        return SC_INTERNAL_SERVER_ERROR;
    }

    private String lookupUserId() throws Exception {
        PreparedQuery pq = getDatastoreService().prepare(new Query("OAuth2Credential").setKeysOnly());
        for (Entity e : pq.asIterable()) {
            if (e.getKey() != null && !isEmpty(e.getKey().getName())) {
                return e.getKey().getName();
            }
        }
        throw new Exception("need user identity");
    }

    private class JobsLoader implements Callable<JobList> {
        public JobList call() throws Exception {
            return bigquery.jobs().list(PROJECT_ID).execute();
        }
    }

    private class JobLoader implements Callable<Job> {
        private final String jobId;

        private JobLoader(String jobId) {
            this.jobId = jobId;
        }

        public Job call() throws Exception {
            return bigquery.jobs().get(PROJECT_ID, this.jobId).execute();
        }
    }

    private class JobCreator implements Callable<Job> {
        private final String query;

        private JobCreator(String query) {
            this.query = query;
        }

        public Job call() throws Exception {
            Job queryJob = makeJob(query);
            return bigquery.jobs().insert(queryJob).setProjectId(PROJECT_ID).execute();
        }

        private Job makeJob(String query) {
            Jobconfigurationquery jobconfigurationquery = new Jobconfigurationquery();

            jobconfigurationquery.setQuery(query);
            jobconfigurationquery.setCreateDisposition("CREATE_IF_NEEDED");

            Jobconfiguration jobconfiguration = new Jobconfiguration();
            jobconfiguration.setQuery(jobconfigurationquery);

            Jobreference jobreference = new Jobreference();
            jobreference.setProjectId(PROJECT_ID);

            Job newJob = new Job();
            newJob.setConfiguration(jobconfiguration);
            newJob.setJobReference(jobreference);

            return newJob;
        }
    }
}
